#include "main.h"

#define MAX_CHAR_DOUBLE 0x7f
#define IS_DEBUG

const int THREAD = 8;
const char *input_path = "./input_image_new.ppm";
const char *check_path = "./check_new.ppm";
const char *output_path = "./output_labels.ppm";

#define square(a) ((a)*(a))
#define max(a, b) (((a)>(b))?(a):(b))
#define min(a, b) (((a)<(b))?(a):(b))
#pragma GCC optimize(3)


int main() {
    // 设置为8线程运行
    omp_set_num_threads(THREAD);
    double first_time, last_time;
#ifdef IS_DEBUG
    double start_time, end_time;
    first_time = omp_get_wtime();
    start_time = omp_get_wtime();
    //读取文件信息
    load();
    end_time = omp_get_wtime();
    printf("Read file spend time: %.3lf s\n", end_time - start_time);
    start_time = omp_get_wtime();
    //转换lab
    cache_rgb2xyz();
    rgb_lab_nocache();
    //计算聚类中心
    find_cluster();
    end_time = omp_get_wtime();
    printf("Find cluster spend time: %.3lf s\n", end_time - start_time);
    start_time = omp_get_wtime();
    //计算超像素
    calculate_super_pixel();
    end_time = omp_get_wtime();
    printf("Calculate super pixel spend time: %.3lf s\n", end_time - start_time);

    start_time = omp_get_wtime();
    //连接孤立像素
    enforce_label_connect();
    end_time = omp_get_wtime();
    printf("Enforce label connect spend time: %.3lf s\n", end_time - start_time);
    start_time = omp_get_wtime();
    //检查结果并写入
    int num = check();
    end_time = omp_get_wtime();
    if (num == 0) {
        printf("Check result and write spend time: %.3lf s\n", end_time - start_time);
    } else {
        printf("Check result spend time: %.3lf s. Error write %d diff\n", end_time - start_time, num);
    }
    start_time = omp_get_wtime();
    printf("There are %d points' labels are different from original file.\n", num);
    //释放资源
    dispose();
    end_time = omp_get_wtime();
    printf("Dispose spend time: %.3lf s\n", end_time - start_time);
    last_time = omp_get_wtime();
    printf("All spend time: %.3lf s\n", last_time - first_time);
#else
    load();
    first_time = omp_get_wtime();
    cache_rgb2xyz();
    rgb_lab_nocache();
    //计算聚类中心
    find_cluster();
    calculate_super_pixel();
    enforce_label_connect();
    last_time = omp_get_wtime();
    printf("Computing time=%.0fms\n", (last_time - first_time) * 1000);
    int num = check();
    printf("There are %d points' labels are different from original file.\n", num);
    dispose();
#endif
    return 0;
}

/**
 * 载入文件
 */
void load() {
    char header[1024];
    int line = 0;
    //以读写的方式打开Filename中的图片文件
    FILE *fp = fopen(input_path, "rb");

    // 读取图像类型，如:P6
    // 从指定的流 fp 读取全部行，并把它存储在 header 所指向的字符串内。
    while (line < 2) {
        fgets(header, 1024, fp);
        if (header[0] != '#') {
            ++line;
        }
    }

    // 读取宽度和高度存入相应的属性中
    sscanf(header, "%d %d\n", &m_width, &m_height);
    m_size = m_height * m_width;

    //初始化所有数据
    initialize();
    // read the maximum of pixels
    fgets(header, 20, fp);
    fread(rgb, m_size * 3, 1, fp);


    // ofc, later, you'll have to cleanup
    //最后一起释放
//    delete[] rgb;
    fclose(fp);
}

void initialize() {
    lab = (double *) malloc(m_size * 3 * sizeof(double));
    rgb = (unsigned char *) malloc(m_size * 3 * sizeof(unsigned char));
    seeds = (double *) malloc(K * 5 * sizeof(double));
    belong = (int *) malloc(m_size * sizeof(int));
    labels = (int *) malloc(m_size * sizeof(int));
    N = K;
}

void cache_rgb2xyz() {
#pragma omp simd
    for (int i = 0; i < 11; ++i) {
        tableRGB[i] = i * (1.0 / 3294.6);
    }
#pragma omp simd
    for (int i = 11; i < 256; ++i) {
        tableRGB[i] = pow((i * (1.0 / 269.025) + 0.0521327014218009), 2.4);
    }
}

void rgb_lab_nocache() {
#pragma omp parallel for default(none) schedule(guided) shared(m_size, rgb, lab)
    for (int i = 0; i < m_size; i++) {
        int k = 3 * i;
        rgb2lab_nocache(rgb + k, lab + k);
//        rgb2lab_vector(rgb + k, lab + k);
    }
}

inline void rgb2lab_nocache(const unsigned char *sRGB, double *lab_val) {
    //------------------------
    // sRGB to XYZ conversion
    //------------------------
    double b0 = tableRGB[sRGB[0]];
    double g0 = tableRGB[sRGB[1]];
    double r0 = tableRGB[sRGB[2]];

    double X, Y, Z;
    X = r0 * 0.4124564 + g0 * 0.3575761 + b0 * 0.1804375;
    Y = r0 * 0.2126729 + g0 * 0.7151522 + b0 * 0.0721750;
    Z = r0 * 0.0193339 + g0 * 0.1191920 + b0 * 0.9503041;

//  此处微调优化浮点数除法
    double fx, fy, fz;
    if (X > 0.008417238336) fx = cbrt(X * (1.0 / 0.950456));
    else fx = (8.192982069151272 * X + 0.1379310344827586);
    if (Y > 0.008856) fy = cbrt(Y);
    else fy = (7.787068965517241 * Y + 0.1379310344827586);
    if (Z > 0.009642005424) fz = cbrt(Z * (1.0 / 1.088754));
    else fz = (7.152275872710678 * Z + 0.1379310344827586);

    lab_val[0] = 200.0 * (fy - fz);
    lab_val[1] = 500.0 * (fx - fy);
    lab_val[2] = 116.0 * fy - 16.0;
}

/**
 * 计算梯度并寻找聚类中心
 */
void find_cluster() {
    step = sqrt(1.0 * m_size / K);
    int offset = (int) (step / 2);

    int n = 0;
    int r = 0;
    int temp = 0;
    for (int y = 0; y < m_height; y++) {
        int Y = (int) (y * step + offset);
        if (Y > m_height - 1) break;

        for (int x = 0; x < m_width; x++) {
            temp++;
            //int X = x*step + xoff;//square grid
            int X = (int) (x * step + (offset << (r & 0x1)));//hex grid
            if (X > m_width - 1) break;

            int i = (Y * m_width + X) * 3;
            seeds[n] = lab[i];
            seeds[n + 1] = lab[i + 1];
            seeds[n + 2] = lab[i + 2];
            seeds[n + 3] = X;
            seeds[n + 4] = Y;
            n += 5;
        }
        r++;
    }
    N = n / 5;
    for (int i = 0; i < N; i++) {
        int m = 5 * i;
        int ox = (int) seeds[m + 3];//原始的x
        int oy = (int) seeds[m + 4];//原始的y
        int oind = oy * m_width + ox;
        //调用确定聚类中心以及周围八个点的梯度最低点
        int storeind = calculate_center_edges(ox, oy);
        if (storeind != oind) {
            seeds[m] = lab[storeind * 3];
            seeds[m + 1] = lab[storeind * 3 + 1];
            seeds[m + 2] = lab[storeind * 3 + 2];
            seeds[m + 3] = storeind % m_width;
            seeds[m + 4] = storeind / m_width;
        }
    }
}

inline int calculate_center_edges(int ox, int oy) {
    const int dx8[8] = {-1, -1, 0, 1, 1, 1, 0, -1};
    const int dy8[8] = {0, -1, -1, -1, 0, 1, 1, 1};
    int oind = oy * m_width + ox;//原始的聚类中心
    for (int j = 0; j < 8; j++) {
        int nx = ox + dx8[j];//新的 x
        int ny = oy + dy8[j];//新的 y
        if (nx >= 0 && nx < m_width && ny >= 0 && ny < m_height) {
            int nind = ny * m_width + nx;
            double x0 = square(lab[oind * 3 - 3] - lab[oind * 3 + 3]) +
                        square(lab[oind * 3 - 2] - lab[oind * 3 + 4]) +
                        square(lab[oind * 3 - 1] - lab[oind * 3 + 5]);
            double y0 = square(lab[oind * 3 - 3 * m_width] - lab[oind * 3 + 3 * m_width]) +
                        square(lab[oind * 3 - 3 * m_width + 1] - lab[oind * 3 + 3 * m_width + 1]) +
                        square(lab[oind * 3 - 3 * m_width + 2] - lab[oind * 3 + 3 * m_width + 2]);

            double x1 = square(lab[nind * 3 - 3] - lab[nind * 3 + 3]) +
                        square(lab[nind * 3 - 2] - lab[nind * 3 + 4]) +
                        square(lab[nind * 3 - 1] - lab[nind * 3 + 5]);

            double y1 = square(lab[nind * 3 - 3 * m_width] - lab[nind * 3 + 3 * m_width]) +
                        square(lab[nind * 3 - 3 * m_width + 1] - lab[nind * 3 + 3 * m_width + 1]) +
                        square(lab[nind * 3 - 3 * m_width + 2] - lab[nind * 3 + 3 * m_width + 2]);
            if ((x1 + y1) < (x0 + y0)) {
                oind = nind;
            }
        }
    }
    return oind;
}

/**
 * 计算超像素
 */
void calculate_super_pixel() {
    const int TIME = 10;//迭代次数
    const int STEP = (int) (sqrt((double) (m_size) / (double) (K)) + 2.0);
    int offset = STEP;
    if (STEP < 10) offset = (int) (STEP * 1.5);

    double inv_xy = 1.0 / (STEP * STEP);
    int width = m_width;
    int height = m_height;

    double max_lab[N] __attribute__((aligned(32)));
    double max_lab_div[N] __attribute__((aligned(32)));

    double sigma[THREAD][N][5] __attribute__((aligned(32)));
    double max_lab_t[THREAD][N] __attribute__((aligned(32)));
    int cluster_size[THREAD][N] __attribute__((aligned(32)));


#pragma omp simd
    for (int i = 0; i < N; ++i) {
        max_lab[i] = 100;
        max_lab_div[i] = 0.01;
    }

    //迭代TIME次
    for (int t = 0; t < TIME; ++t) {
        //初始化清零数据
        memset(cluster_size, 0, N * THREAD * sizeof(int));
        memset(max_lab_t, 0, N * THREAD * sizeof(double));
        memset(sigma, 0, 5 * N * THREAD * sizeof(double));
        //遍历所有seed,计算最近的靠近位置
#pragma omp parallel for default(none) shared(height, width, N, offset, t, seeds, lab, belong, inv_xy, max_lab_t, max_lab_div, sigma, cluster_size) schedule(guided, 10)
        for (int y = 0; y < height; ++y) {
            double dis_vec_y[width] __attribute__((aligned(32)));
            double dis_lab_y[width] __attribute__((aligned(32)));
            memset(dis_vec_y, MAX_CHAR_DOUBLE, width * sizeof(double));
            const int y_index = y * width;
            for (int n = 0; n < N; n++) {
                double *seed = seeds + 5 * n;
                if ((int) (seed[4] - offset) <= y && y < (int) (seed[4] + offset)) {
                    const int x1 = max(0, (int) (seed[3] - offset));
                    const int x2 = min(width, (int) (seed[3] + offset));
                    const double div_lab = max_lab_div[n];
                    for (int x = x1; x < x2; ++x) {
                        int i = y_index + x;
                        dis_lab_y[x] = square(lab[i * 3] - seed[0]) +
                                       square(lab[i * 3 + 1] - seed[1]) +
                                       square(lab[i * 3 + 2] - seed[2]);
                    }
                    for (int x = x1; x < x2; ++x) {
                        int i = y_index + x;
                        double temp_dis_xy = square(x - seed[3]) + square(y - seed[4]);
                        double dist = dis_lab_y[x] * div_lab + temp_dis_xy * inv_xy;
                        if (dist < dis_vec_y[x]) {
                            dis_vec_y[x] = dist;
                            belong[i] = n;
                        }
                    }
                }
            }
            const int thread_num = omp_get_thread_num();
            for (int x = 0; x < width; ++x) {
                int i = width * y + x;
                int k = belong[i];
                if (max_lab_t[thread_num][k] < dis_lab_y[x]) {
                    max_lab_t[thread_num][k] = dis_lab_y[x];
                }
#pragma omp simd
                for (int j = 0; j < 3; ++j) {
                    sigma[thread_num][k][j] += lab[i * 3 + j];
                }
                sigma[thread_num][k][3] += x;
                sigma[thread_num][k][4] += y;
                cluster_size[thread_num][k]++;
            }
        }

        // 重新计算种子点
        for (int k = 0; k < N; k++) {
            int seed_size = 0;
            double sigma_t[5] __attribute__((aligned(32))) = {0};
            for (int i = 0; i < THREAD; ++i) {
#pragma omp simd
                for (int j = 0; j < 5; ++j) {
                    sigma_t[j] += sigma[i][k][j];
                }
                if (max_lab[k] < max_lab_t[i][k]) {
                    max_lab[k] = max_lab_t[i][k];
                }
                seed_size += cluster_size[i][k];
            }
            if (seed_size == 0) seed_size = 1;
            double inv = 1.0 / seed_size;
            max_lab_div[k] = 1.0 / max_lab[k];
#pragma omp simd
            for (int i = 0; i < 5; ++i) {
                seeds[5 * k + i] = sigma_t[i] * inv;
            }
        }
    }
}

void enforce_label_connect() {
//	const int dx8[8] = {-1, -1,  0,  1, 1, 1, 0, -1};
//	const int dy8[8] = { 0, -1, -1, -1, 0, 1, 1,  1};

    int dx4[4] = {-1, 0, 1, 0};
    int dy4[4] = {0, -1, 0, 1};

    int sz = m_size;
    const int THRESHOLD = (m_size / K) >> 2;
    int width = m_width;
    int height = m_height;

    int *tempLabel = (int *) malloc(sz * sizeof(int));
    memset(tempLabel, -1, m_size * sizeof(int));
    memset(labels, -1, m_size * sizeof(int));

    int thread0 = min(omp_get_max_threads(), 8);
    int thread1 = min(omp_get_max_threads(), 16);

    struct list P[N * N];
#pragma omp simd
    for (int i = 0; i < N * N; ++i) {
        struct list temp = {16, 0, malloc(16 * sizeof(int))};
        P[i] = temp;
    }
    //vector<int> P[sz];

//    vector<int> G[THREAD][N];
    struct list Q[THREAD][N];
#pragma omp simd
    for (int i = 0; i < THREAD; ++i) {
        for (int j = 0; j < N; ++j) {
            struct list temp = {16, 0, malloc(16 * sizeof(int))};
            Q[i][j] = temp;
        }
    }

#pragma omp parallel for num_threads(thread0) default(none) shared(sz, Q, belong, N)
    for (int i = 0; i < sz; i++) {
        int tid = omp_get_thread_num();
        add(&Q[tid][belong[i]], i);
    }

    int label = 0;
    int adjlabel = 0;//adjacent label
    int oindex = 0;

    int thread_num = THREAD;
    int K0 = K;

#pragma omp parallel for num_threads(thread1) default(none) shared(thread_num, K0, N, Q, tempLabel, oindex, width, dx4, dy4, height, belong, P)
    for (int id = 0; id < N; id++) {
        int nowLabel = id;

        int size = 0;
        for (int tid = 0; tid < thread_num; tid++)
            size += Q[tid][id].length;
        int isOK = 0;
        int *arr = (int *) malloc(size * sizeof(int));
        while (isOK < size) {
            int now = -1;
            for (int tid = 0; tid < thread_num; tid++) {
                for (int i = 0; i < Q[tid][id].length; i++) {
                    if (tempLabel[get(Q[tid][id], i)] == -1) {
                        now = get(Q[tid][id], i);
                        break;
                    }
                }
            }
            int start = 0, finish = 0;
            arr[finish++] = now;
            tempLabel[now] = nowLabel;
            oindex = now;
            while (start < finish) {
                int k = arr[start++];
                int x = k % width;
                int y = k / width;
                for (int i = 0; i < 4; i++) {
                    int xx = x + dx4[i];
                    int yy = y + dy4[i];
                    if ((xx >= 0 && xx < width) && (yy >= 0 && yy < height)) {
                        int nindex = yy * width + xx;

                        if (0 > tempLabel[nindex] && belong[oindex] == belong[nindex]) {
                            arr[finish++] = nindex;
                            tempLabel[nindex] = nowLabel;
                        }
                    }
                }

            }
            cut(&P[nowLabel], finish);
            for (int i = 0; i < finish; i++)
                set(P[nowLabel], i, arr[i]);
            isOK += finish;
            nowLabel += K0;
        }

        free(arr);
    }

    oindex = 0;
    for (int j = 0; j < height; j++) {
        for (int k = 0; k < width; k++) {
            if (0 > labels[oindex]) {
                labels[oindex] = label;
                int bel = tempLabel[oindex];
                int count2 = P[bel].length;
                if (count2 <= THRESHOLD) {
                    for (int n = 3; n >= 0; n--) {
                        int x = k + dx4[n];
                        int y = j + dy4[n];
                        if ((x >= 0 && x < width) && (y >= 0 && y < height)) {
                            int nindex = y * width + x;
                            if (labels[nindex] >= 0) {
                                adjlabel = labels[nindex];
                                break;
                            }
                        }
                    }

#pragma omp parallel for num_threads(thread0) default(none) shared(count2, labels, P, bel, adjlabel)
                    for (int c = 0; c < count2; c++) {
                        labels[get(P[bel], c)] = adjlabel;
                    }
                    label--;
                } else {
#pragma omp parallel for num_threads(thread0) default(none) shared(count2, labels, P, bel, label)
                    for (int c = 0; c < count2; c++) {
                        labels[get(P[bel], c)] = label;
                    }
                }
                label++;
            }
            oindex++;
        }
    }
    free(tempLabel);

#pragma omp simd
    for (int i = 0; i < N * N; ++i) {
        dispose_list(P[i]);
    }
#pragma omp simd
    for (int i = 0; i < THREAD; ++i) {
        for (int j = 0; j < N; ++j) {
            dispose_list(Q[i][j]);
        }
    }
}

//inline void rgb2lab_vector(const unsigned char *sRGB, double *lab_val) {
//    //------------------------
//    // sRGB to XYZ conversion
//    //------------------------
//    __m256d r256 = _mm256_set1_pd(tableRGB[sRGB[2]]);
//    __m256d g256 = _mm256_set1_pd(tableRGB[sRGB[1]]);
//    __m256d b256 = _mm256_set1_pd(tableRGB[sRGB[0]]);
//    __m256d xyz = _mm256_mul_pd(r256, RGB_CONST[0]);
//    xyz = _mm256_fmadd_pd(g256, RGB_CONST[1], xyz);
//    xyz = _mm256_fmadd_pd(b256, RGB_CONST[2], xyz);
//    xyz = _mm256_mul_pd(xyz, RGB_CONST[3]);
//    // 向量化三元运算符
//    xyz = _mm256_blendv_pd(_ZGVdN4vv_pow(xyz, RGB_CONST[7]),
//                           _mm256_fmadd_pd(xyz, RGB_CONST[4], RGB_CONST[5]),
//                           _mm256_cmp_pd(xyz, RGB_CONST[6], _CMP_LE_OQ));
//    double xyz0[4] __attribute__((aligned(32)));
//    _mm256_stream_pd(xyz0, xyz);
//    double X, Y, Z;
//    X = xyz0[0];
//    Y = xyz0[1];
//    Z = xyz0[2];
//
//    lab_val[2] = 116.0 * Y - 16.0;
//    lab_val[1] = 500.0 * (X - Y);
//    lab_val[0] = 200.0 * (Y - Z);
//}

int check() {
    char header[1024];
    int line = 0;
    int num = 0;
//    int sub_num = 0;

    FILE *fp = fopen(check_path, "rb");

    // read the image type, such as: P6
    // skip the comment lines
    while (line < 2) {
        fgets(header, 1024, fp);
        if (header[0] != '#') {
            ++line;
        }
    }
    // read width and height
    int w = 0;
    int h = 0;
    sscanf(header, "%d %d\n", &w, &h);
    if (w != m_width || h != m_height) return -1;

    // read the maximum of pixels
    fgets(header, 20, fp);

    // get check_rgb data
    unsigned char *check_rgb = malloc(m_size * 3 * sizeof(unsigned char));

    fread(check_rgb, m_size * 3, 1, fp);

#pragma omp parallel for default(none) schedule(guided) shared(m_size, check_rgb, labels, num)
    for (int i = 0; i < m_size; i++) {
        int k = 3 * i;
        int ground;
        ground = check_rgb[k + 2] << 16; // r
        ground |= check_rgb[k + 1] << 8;  // g
        ground |= check_rgb[k + 0];       // b
        if (ground != labels[i]) {
#pragma omp atomic
            num++;
        }
    }
    fclose(fp);

    fp = fopen(output_path, "wb");
    //没有不一样的像素点,直接写出
    if (num == 0) {
        // write the PPM header info, such as type, width, height and maximum
        fprintf(fp, "P6\n%d %d\n255\n", m_width, m_height);
        fwrite(check_rgb, m_size * 3, 1, fp);
    } else {
        // write the PPM header info, such as type, width, height and maximum
        fprintf(fp, "P6\n%d %d\n255\n", m_width, m_height);
        unsigned char *rgbw = malloc(m_size * 3 * sizeof(unsigned char));
        int k = 0;
        for (int i = 0; i < m_height; i++) {
            for (int j = 0; j < m_width; j++) {
                rgbw[i * m_width * 3 + j * 3 + 2] = labels[k] >> 16 & 0xff;  // r
                rgbw[i * m_width * 3 + j * 3 + 1] = labels[k] >> 8 & 0xff;  // g
                rgbw[i * m_width * 3 + j * 3 + 0] = labels[k] & 0xff;  // b

                // rgbw[i*(width) + j + 0] = c;
                k++;
            }
        }
        //最后一起释放
        free(rgbw);
    }
    fclose(fp);
    free(check_rgb);
    return num;
}

void dispose() {
    free(rgb);
    free(labels);
    free(lab);
    free(seeds);
    free(belong);
}

// 模拟vector操作
inline void add(struct list *l, int a) {
    if (l->length == l->max_size - 1) {
        int *t = l->array;
        int new_size = l->max_size << 2;
        l->array = malloc(sizeof(int) * new_size);
        memcpy(l->array, t, l->max_size * sizeof(int));
        l->max_size = new_size;
        free(t);
    }
    l->array[l->length] = a;
    l->length++;
}

inline void set(struct list l, int i, int a) {
    l.array[i] = a;
}

inline int get(struct list l, int i) {
    return l.array[i];
}

inline void dispose_list(struct list l) {
    if (l.array) {
        free(l.array);
    }
}

inline void cut(struct list *l, int i) {
    int *res = (int *) malloc(i * sizeof(int));
    if (l->length > i){
        memcpy(res, l->array, i * sizeof(int));
    } else {
        memcpy(res, l->array, l->length * sizeof(int));
    }
    free(l->array);
    l->array = res;
    l->length = i;
}